import cv2
import pytesseract
import sys
import re
from pathlib import Path
from aip import AipOcr
from contract_renamer.config import load_config
from PIL import Image
import logging
import os

class OCRProcessor:
    def __init__(self):
        config = load_config()
        baidu_config = config.get('baidu_ocr', {})
        app_id = baidu_config.get('app_id', '')
        api_key = baidu_config.get('api_key', '')
        secret_key = baidu_config.get('secret_key', '')
        
        # 确保 app_id 是字符串类型
        self.baidu_client = AipOcr(
            str(app_id) if app_id else '',
            api_key,
            secret_key
        )
        
        # 设置 Tesseract 路径
        if getattr(sys, 'frozen', False):
            # 打包后的路径
            base_path = sys._MEIPASS if hasattr(sys, '_MEIPASS') else os.path.dirname(sys.executable)
            tesseract_path = os.path.join(base_path, 'Resources', 'tesseract', 'tesseract')
            if os.path.exists(tesseract_path):
                pytesseract.pytesseract.tesseract_cmd = tesseract_path
        elif sys.platform.startswith('win'):
            # Windows 开发环境
            pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

    def baidu_ocr(self, image_path):
        """使用百度 OCR API 识别图片"""
        try:
            # 检查 API 配置
            config = load_config()
            baidu_config = config.get('baidu_ocr', {})
            app_id = baidu_config.get('app_id')
            api_key = baidu_config.get('api_key')
            secret_key = baidu_config.get('secret_key')

            if not all([app_id, api_key, secret_key]):
                error_msg = "百度 OCR API 未配置，请先配置API密钥"
                self.log_error(error_msg)
                return None, error_msg

            # 重新初始化百度OCR客户端（确保使用最新配置）
            self.baidu_client = AipOcr(str(app_id), api_key, secret_key)

            # 读取图片文件
            with open(image_path, 'rb') as f:
                image = f.read()
            
            # 调用百度OCR API
            result = self.baidu_client.basicGeneral(image)
            
            # 检查错误码
            if 'error_code' in result:
                error_msg = f"百度OCR API 调用失败: {result.get('error_msg', '未知错误')}"
                self.log_error(error_msg)
                return None, error_msg
            
            # 提取识别文本
            if 'words_result' in result:
                text = '\n'.join(item['words'] for item in result['words_result'])
                return text, None
            return None, "百度OCR未返回识别结果"

        except Exception as e:
            error_msg = f"百度OCR处理失败: {str(e)}"
            self.log_error(error_msg)
            return None, error_msg

    def tesseract_ocr(self, image_path):
        """使用 Tesseract 进行 OCR 识别"""
        try:
            img = cv2.imread(str(image_path))
            if img is None:
                self.log_error(f"无法读取图片: {image_path}")
                return None

            # 图像预处理
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            text = pytesseract.image_to_string(gray, lang='chi_sim')
            
            # 记录识别结果
            if text:
                self.log_recognition(f"Tesseract识别结果:\n{text}\n---")
            else:
                self.log_error("Tesseract未能识别出文本")
            
            return text
        except Exception as e:
            self.log_error(f"Tesseract处理失败: {str(e)}")
            return None

    def process_text(self, text):
        """处理识别出的文本，提取文档类型和合同号"""
        if not text:
            return None, None, None

        # 记录原始识别文本
        self.log_recognition("OCR原始文本:\n" + text + "\n---")

        # 识别文档类型
        doc_type = None
        if any(exact_match in text for exact_match in ["到货证明", "到贷证明", "到作证明"]):
            doc_type = "到货证明"
        elif "验收证书" in text or "验收证明" in text:
            doc_type = "验收证明"

        # 识别合同号
        contract_number = None
        # 合同号匹配模式
        patterns = [
            # 标准格式
            r'合同[编号号].?\s*[：:\s]*([XKx][Jj][Ss]\d{4,})',
            r'编号\s*[：:\s]*([XKx][Jj][Ss]\d{4,})',
            
            # 宽松格式（允许更多的分隔符和空格）
            r'合同.{0,4}号.{0,4}([XKx][Jj][Ss]\d{4,})',
            r'([XKx][Jj][Ss]\d{4,})',
            
            # 特殊格式
            r'合同编号[:：\s]*(\d{4}[XKx][Jj][Ss]\d{4,})',
            r'(\d{4}[XKx][Jj][Ss]\d{4,})',
            
            # 处理X的误识别
            r'[Xx×Ⅹ][Jj][Ss]\d{4,}',  # X可能被识别为×或Ⅹ
            r'[kK][Jj][Ss]\d{4,}',     # X可能被识别为K
            
            # 处理J的误识别（不包括数字1）
            r'[Xx][JjIilL][Ss]\d{4,}',  # J可能被识别为I,i,l,L
            r'[Xx][Tt][Ss]\d{4,}',      # J可能被识别为T
            
            # 处理S的误识别
            r'[Xx][Jj][Ss§]\d{4,}',     # S可能被识别为§
            r'[Xx][Jj][Cc]\d{4,}',      # S可能被识别为C
            
            # 处理数字的误识别
            r'[Xx][Jj][Ss][oO0-9]+',    # 数字中的0可能被识别为o或O
            r'[Xx][Jj][Ss]\d+',         # 标准数字格式
            
            # 组合错误处理
            r'[XxKk×Ⅹ][JjIilLTt][Ss§Cc][0-9oO]+',
            
            # 处理空格和特殊字符
            r'[Xx]\s*[Jj]\s*[Ss]\s*\d{4,}',
            r'[Xx][Jj][Ss][-_]\d{4,}'
        ]

        # 尝试所有模式匹配
        text = text.upper()  # 转换为大写以统一处理
        for pattern in patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                try:
                    contract_number = match.group(1).upper()
                except IndexError:
                    contract_number = match.group(0).upper()
                
                # 清理合同号中的特殊字符
                contract_number = re.sub(r'[^A-Z0-9]', '', contract_number)
                
                # 修正常见的OCR错误
                ocr_corrections = {
                    'K': 'X',    # K -> X
                    '×': 'X',    # × -> X
                    'Ⅹ': 'X',    # Ⅹ -> X
                    'I': 'J',    # I -> J
                    'L': 'J',    # L -> J
                    'T': 'J',    # T -> J
                    'l': 'J',    # l -> J
                    '§': 'S',    # § -> S
                    'C': 'S',    # C -> S
                    'O': '0',    # O -> 0
                    'o': '0',    # o -> 0
                    # 移除 '1': 'J' 的转换
                }
                
                for wrong, correct in ocr_corrections.items():
                    contract_number = contract_number.replace(wrong, correct)
                
                # 确保合同号符合基本格式
                if re.match(r'^.*?XJS\d+$', contract_number):
                    # 如果合同号前面有其他字符，只保留XJS及后面的数字
                    if not contract_number.startswith('XJS'):
                        match = re.search(r'XJS\d+', contract_number)
                        if match:
                            contract_number = match.group(0)
                    break

        # 记录识别结果
        if doc_type or contract_number:
            self.log_recognition(f"识别结果 - 文档类型: {doc_type}, 合同号: {contract_number}")
        else:
            self.log_recognition("未能识别文档类型和合同号")
        self.log_recognition("---\n")

        return doc_type, contract_number, text

    def log_error(self, message):
        """记录错误日志"""
        log_dir = Path.home() / '.contract_renamer' / 'logs'
        log_dir.mkdir(parents=True, exist_ok=True)
        log_file = log_dir / 'error.log'
        
        from datetime import datetime
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"[{timestamp}] {message}\n")

    def log_recognition(self, message):
        """记录识别结果到日志"""
        log_dir = Path.home() / '.contract_renamer' / 'logs'
        log_dir.mkdir(parents=True, exist_ok=True)
        log_file = log_dir / 'recognition.log'
        
        from datetime import datetime
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"[{timestamp}] {message}\n")

class OCRUtils:
    def recognize_text_tesseract(self, image_path):
        """使用Tesseract进行OCR识别"""
        try:
            image = Image.open(image_path)
            text = pytesseract.image_to_string(image, lang='chi_sim')
            return text.strip()
        except Exception as e:
            logging.error(f"Tesseract OCR失败: {str(e)}")
            return ""

    def recognize_text_baidu(self, image_path):
        """使用百度OCR进行识别"""
        try:
            with open(image_path, 'rb') as fp:
                image = fp.read()
            
            result = self.client.basicGeneral(image)
            
            if 'words_result' in result:
                text = '\n'.join([w['words'] for w in result['words_result']])
                return text.strip()
            return ""
        except Exception as e:
            logging.error(f"百度OCR失败: {str(e)}")
            return ""